import torch
import torch.nn as nn
import torch.nn.functional as F
from .base import *


# class FPN(nn.Module):
#     def __init__(self,in_channels_list,out_channels):
#         super(FPN,self).__init__()
#         leaky = 0
#         if (out_channels <= 64):
#             leaky = 0.1
#         self.output1 = conv_bn1X1(in_channels_list[0], out_channels, stride = 1, leaky = leaky)
#         self.output2 = conv_bn1X1(in_channels_list[1], out_channels, stride = 1, leaky = leaky)
#         self.output3 = conv_bn1X1(in_channels_list[2], out_channels, stride = 1, leaky = leaky)

#         self.merge1 = conv_bn(out_channels, out_channels, leaky = leaky)
#         self.merge2 = conv_bn(out_channels, out_channels, leaky = leaky)

#     def forward(self, input):
#         input = list(input.values())

#         output1 = self.output1(input[0])
#         output2 = self.output2(input[1])
#         output3 = self.output3(input[2])


#         up3 = F.interpolate(output3, size=[output2.size(2), output2.size(3)], mode="nearest")
#         output2 = output2 + up3
#         output2 = self.merge2(output2)

#         up2 = F.interpolate(output2, size=[output1.size(2), output1.size(3)], mode="nearest")
#         output1 = output1 + up2
#         output1 = self.merge1(output1)

#         # out = [output1, output2, output3]
#         # return out

#         return output1



class FPN(nn.Module):
    def __init__(self, in_channels_list, out_channels, in_size=(640, 640)):
        super(FPN,self).__init__()
        self.output1 = conv_bn1X1(in_channels_list[0], out_channels, stride = 1)
        self.output2 = conv_bn1X1(in_channels_list[1], out_channels, stride = 1)
        self.output3 = conv_bn1X1(in_channels_list[2], out_channels, stride = 1)
        self.output4 = conv_bn1X1(in_channels_list[3], out_channels, stride = 1)

        self.merge1 = conv_bn(out_channels, out_channels)
        self.merge2 = conv_bn(out_channels, out_channels)
        self.merge3 = conv_bn(out_channels, out_channels)
        self.in_size = in_size

    def forward(self, input):
        output1 = self.output1(input[0])
        output2 = self.output2(input[1])
        output3 = self.output3(input[2])
        output4 = self.output4(input[3])


        # up4 = F.interpolate(output4, size=[i//16 for i in self.in_size], mode="nearest")
        up4 = F.interpolate(output4, scale_factor=2, mode="nearest")
        output3 = output3 + up4
        output3 = self.merge3(output3)

        # up3 = F.interpolate(output3, size=[i//8 for i in self.in_size], mode="nearest")
        up3 = F.interpolate(output3, scale_factor=2, mode="nearest")
        output2 = output2 + up3
        output2 = self.merge2(output2)


        # up2 = F.interpolate(output2, size=[i//4 for i in self.in_size], mode="nearest")
        up2 = F.interpolate(output2, scale_factor=2, mode="nearest")
        output1 = output1 + up2
        output1 = self.merge1(output1)
        return output1